import torch
import torch.nn.functional as F

def weighted_average(tensor, weight, dim=1):
    tensor = tensor * weight
    return torch.sum(tensor, dim=dim) / torch.sum(weight, dim=dim)

def index_only_dist(result, target, index, mask):
    n_dim = 1
    v_dim = 2
    # [B N]
    return weighted_average(
        torch.sum(
            F.relu(
                F.log_softmax(target, dim=v_dim) - F.log_softmax(result, dim=v_dim)
            ) * index, dim=v_dim
        ),
        mask,
        dim=n_dim
    )

def kl_divergence_dist(result, target, index, mask):
    n_dim = 1
    v_dim = 2
    return weighted_average(
        torch.sum(
            F.softmax(target, dim=v_dim) * ( - F.log_softmax(result, dim=v_dim) + 
            F.log_softmax(target, dim=v_dim)),
            dim=v_dim
        ),
        mask,
        dim=n_dim
    )

def js_divergence_dist(result, target, index, mask):
    return 0.5 * (kl_divergence_dist(result, target, index, mask) +
             kl_divergence_dist(target, result, index, mask))

def cross_entropy_dist(result, target, index, mask):
    n_dim = 1
    v_dim = 2
    return weighted_average(
        torch.sum(
            - F.softmax(target, dim=v_dim) * F.log_softmax(result, dim=v_dim),
            dim=v_dim
        ),
        mask,
        dim=n_dim
    )